#/*************************************************************************
#    > File Name: Lib/Parallel_Operation.py
#    > Author: Yan Wang
#    > Mail: wangyan@imnu.edu.cn
#    > Created Time: 2022年05月26日 星期四 15时23分41秒
# ************************************************************************/
#!/usr/bin/env python
# -*- coding: utf-8 -*-
from Library import *

def PBS_Running(param,pbs_template="./examples/test_input/pbs_template.pbs"):
	print("in PBS_Running, the input param is ", param )
	for folder in param["folder"]:
		pbs_file=folder+"/run_pbs.pbs"
		Copy_File(pbs_template,pbs_file)
		print("pbs_file",pbs_file)
		linenos=[2,7,8,18]
		strings=["#PBS -N {}".format(folder), 
				"#PBS -e {}/job.error".format(folder),
				"#PBS -o {}/job.log".format(folder),
				"cd {}".format(folder)]
		File_Replace(pbs_file,linenos,strings)
		os.system(param["command"].format(folder))
	return()


def Slurm_Running(param,slurm_template="./examples/test_input/slurm_template.pbs"):
	print("in Slurm_Running, the input param is ", param )
	for folder in param["folder"]:
		slurm_file=folder+"/run_slurm.pbs"
		Copy_File(slurm_template,slurm_file)
		print("slurm_file",slurm_file)
		linenos=[2,10]
		strings=["#SBATCH -J {}".format(folder), "cd {}".format(folder)]
		File_Replace(slurm_file,linenos,strings)
		os.system(param["command"].format(folder))
	return()


# 进入多个核，并行计算
# 第一个参数为需要运行的函数， 第二个paras为这个函数的参数组成的列表，  paras中每一个元素会开辟一个新核，运行run_func(paras_element)
def Parallel_Running(run_func, paras, record_file="None",log_file="../log/log.dat"):
	f=open(log_file,'w')
	start_t = datetime.datetime.now()   #.datetime.now()获取指定日期格式的系统的本地时间
	num_cores = int(mp.cpu_count())
	# 使用的核数，不必超过需要的参数个数
	if num_cores> len(paras):
	    num_cores=len(paras)
	print("local PC has " + str(num_cores) + " cores.")
	f.write("This running use " + str(num_cores) + " CPU cores. \n");
	f.write("The start time is : " + str(start_t) + "\n");

	pool = mp.Pool(num_cores)
	
	#print("The input Function is: ", run_func.__name__)
	#print("The input params are: \n", paras.head())
	
	#[print(row.values[0], type(row.values[0])) for index, row in paras.iterrows()]   
	#results=[pool.apply_async(run_func, args = (row.values[0],)) for index, row in paras.iterrows()]   
	results=[]
	for index, row in paras.iterrows():
	    if index%1000 == 0:
	        f.write("   already finished " +  str(index) + "points. \n");
	    end_t=datetime.datetime.now()
	    elapsed_sec=(end_t - start_t).total_seconds()
	    current_time = elapsed_sec%(60*60)
	    if current_time == 0:
	        f.write("   already running " +  str(current_time) + "hours. \n");
	    results.append(pool.apply_async(run_func, args = (row.values[0],)))
    #pool.apply_async()结果的顺序不能保证与调用的顺序相同
	print("results type ***************",type(results))
	if results is None:
		print("Finish Parallel Running once, and no return value.")
	else:
		#results=[print(p) for p in results]
		results=[print(p.get()) for p in results]
	
    
####if record_file != "None":
####	for index_c, content in enumerate(results):
####	    this_file=open(record_file.format(index_c),"w")
####	    this_file.write(content)
####	    this_file.close()
    
	end_t=datetime.datetime.now()
	elapsed_sec=(end_t - start_t).total_seconds()
	print("Multi processes calculation,  total costs are: " +
	        "{:.2f}".format(elapsed_sec) + " second.")
	f.write("The end time is : " + str(end_t) + ". \n Totally use " + str(elapsed_sec) + " seconds. \n");
	return()

# 通过SSH执行命令
# 这些节点最好是添加过 host 的授信节点， 不需要再次填写密码
# 对于已授信节点 ssh_para 应存在 "node_name" 的key， 通过name连接ssh
# 如果需要填写密码，则 ssh_para中 不可存在 node_name的key， 通过ip， user_name, passwd 连接ssh
# 命令应存放在 ssh_para 的 "command" key 中
def Run_SSH(ssh_para):
	ssh = paramiko.SSHClient()   #远程连接服务器，建立一个sshclient对象
	# 允许链接不在 know_hosts文件中的主机
	ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
	# 建立连接
	if "node_name" in ssh_para.keys():
		ssh.connect(ssh_para["node_name"], port=22)
	else:
		ssh.connect(ssh_para["ip"], username=ssh_para["user"], port=22, password=ssh_para["passwd"])
	# 使用这个连接执行命令
	sh_stdin, ssh_stdout, ssh_stderr = ssh.exec_command(ssh_para["command"])
	# 获取输出
	print(ssh_stdout.read())
	# 关闭连接
	ssh.close()

# 使用并行计算的方式，每个核打开一个SSH 连接到集群的一个节点，  每个节点执行一个命令
# 这里command 第一个参数为需要执行的命令， 第二个参数为执行命令使用的节点名称
def Run_Command_with_SSH(ssh_para):

	if "node_name" in ssh_para.keys():
		print(" We will execute the commmand:",ssh_para["command"], "\n in node:",ssh_para["node_name"])
	else:
		print(" We will execute the commmand:",ssh_para["command"], "\n in node:",ssh_para["ip"])

	result = subprocess.run(Run_SSH(ssh_para), capture_output=True, shell=True)  # capture_output=True获stdout和stderr，调用时内部的Popen对象将自动使用stdout=PIPE和stderr = PIPE创建标准输出和标准错误对象；
	return(result)

